
# -*- coding: utf-8 -*-

import os
import sys
import threading
import time
import getopt
import multiprocessing

class MultiRunCMDThread(threading.Thread):
    ''' The Thread Class that will run a CMD

        '''
    def __init__(self, name, cmd):
        threading.Thread.__init__(self)
        self.name=name
        self.cmd=cmd
    def run(self):
        ''' Run the CMD in the thread
            '''
        print 'Thread',self.name,'is running.'
        os.system(' '.join(self.cmd))



import logging;

#-------------------------------------------------------------------------------
def loggingDemo():
    """Just demo basic usage of logging module
    """
    logging.info("You should see this info both in log file and cmd window");
    logging.warning("You should see this warning both in log file and cmd window");
    logging.error("You should see this error both in log file and cmd window");

    logging.debug("You should ONLY see this debug in log file");
    return;

#-------------------------------------------------------------------------------
def initLogging(logFilename):
    """Init for logging
    """
    logging.basicConfig(
                    level    = logging.DEBUG,
                    format   = 'LINE %(lineno)-4d  %(levelname)-8s %(message)s',
                    datefmt  = '%m-%d %H:%M',
                    filename = logFilename,
                    filemode = 'w');
    # define a Handler which writes INFO messages or higher to the sys.stderr
    console = logging.StreamHandler();
    console.setLevel(logging.INFO);
    # set a format which is simpler for console use
    formatter = logging.Formatter('LINE %(lineno)-4d : %(levelname)-8s %(message)s');
    # tell the handler to use this format
    console.setFormatter(formatter);
    logging.getLogger('').addHandler(console);


def help_info():
    print("help infomation");
    print("-r, --root:    the root folder of the project, it should have the tree structure like this:");
    print("--root");
    print("  --VRM");
    print("  --RD");
    print("  --seq");
    print("-f, --flag:    the flag for mode, the first bit indicates projection, the second indicates encoder, \
          the third indicates SPSNR")
    print("  1: projection via VRM")
    print("  2: encode sequence with RD")
    print("  4: spsnr with VRM")


def process_all(root_path, flag):
    '''
    script for projectin and encoder the 360 videos
    '''
    # pylog
    VRM_RD_log = "VRM_RD.log";
    initLogging(VRM_RD_log);
    logging.info(time.strftime("%Y-%m-%d %X", time.localtime()));
    logging.info("projection: VRM, coding:RD with VR extension");

    #flag
    VRM_flag = bool(flag&(0x01 << 0));
    RD_flag = bool(flag&(0x01 << 1));
    SPSNR_flag = bool(flag&(0x01 << 2));
    #PATH
    if root_path == '':
        root_path = os.getcwd();

    VRM_path = r"VRM/bin/"
    SEQ_path = r"seq/"
    RD_ENCODER_path = r"RD/RD16.1_with_vr_extension/source/bin/"
    MAP_path = RD_ENCODER_path;
    REC_path = RD_ENCODER_path;
    AVS_path = RD_ENCODER_path;
    # encoder para
    encoder_para = r"-f encoder_ra.cfg -f encoder.cfg"
    VRM_para = r"-c encoder_ra.cfg -C encoder_ra.cfg -w 3840 -h 1920 -x 3840 -y 1920 -N 30 -M map -F eqar -f eqar"

    # EXE
    VRM=r"VRM.exe"
    ENCODER = r"lencod.exe"

    #log
    VRM_log = r"VRM_py.log";
    RD_log = r"RD_py.log";

    #MAP
    map_yuv_width = r"3840";
    map_yuv_height = r"1920";

    #ext
    map_file_ext=r"_map.yuv";
    src_file_ext=r".yuv";
    rec_file_ext=r"_rec.yuv";
    avs_file_ext=r".avs";

    file_name_all = [[], [], [], []]; #src, map, rec, avs
    file_name_idx = {r"src":0, r"map":1, r"rec":2, r"avs":3};
    QP_list = [27, 32, 38, 45];
    #prepare the file name and seq_temp floder
    seq_temp_path = r"seq_temp";
    if not os.path.exists(seq_temp_path):
        os.mkdir(seq_temp_path);

    for src_file_name in os.listdir(os.path.join(root_path, SEQ_path)):
        name, ext = os.path.splitext(src_file_name);
        if ext == src_file_ext:
            file_name_all[0].append(src_file_name);
            file_name_all[1].append(name+map_file_ext);
            file_name_all[2].append([name+str(QP_list[i])+rec_file_ext for i in range(4)]);
            file_name_all[3].append(name+avs_file_ext);

    #VRM
    if(VRM_flag):
        print(r"VRM...")
        VRM_bin = os.path.join(root_path,VRM_path);
        os.chdir(VRM_bin);
        print(os.getcwd());
        if os.path.isfile(os.path.join(VRM_bin, VRM)):
            src_file_path = os.path.join(root_path, SEQ_path);
            file_name_idx_src = file_name_idx[r"src"];
            for src_file_name_idx in range(len(file_name_all[file_name_idx_src])):
                src_file_name = file_name_all[file_name_idx_src][src_file_name_idx];
                if src_file_name[-len(src_file_ext):] == src_file_ext:
                    print(r"process yuv file:"+src_file_name);
                    file_name_idx_map = file_name_idx[r"map"];
                    output_name = file_name_all[file_name_idx_map][src_file_name_idx];
                    if os.path.isfile(os.path.join(MAP_path, output_name)):
                        print(os.path.join(MAP_path, output_name) + r" exist, delete it");
                        os.remove(os.path.join(MAP_path, output_name));
                    command = os.path.join(VRM_bin, VRM)+r" "+VRM_para+r" -i "+os.path.join(src_file_path,src_file_name) \
                              + " -o " + os.path.join(root_path, os.path.join(MAP_path, output_name))+ r">>"+VRM_log;
                    os.system(r"echo VRM start...");
                    os.system(command);
        print(r"VRM end");

    #RD
    if(RD_flag):
        #encoder
        print(r"RD lencod: encode the seq with QP="+str(QP_list)+"...")
        RD_bin = os.path.join(root_path,RD_ENCODER_path);
        os.chdir(RD_bin);
        print(r"cur_path:"+os.getcwd());
        if os.path.isfile(os.path.join(RD_bin,ENCODER)):
            file_name_idx_map = file_name_idx[r"map"];
            file_name_idx_rec = file_name_idx[r"rec"];
            for map_file_name_idx in range(len(file_name_all[file_name_idx_map])):
                map_file_name = file_name_all[file_name_idx_map][map_file_name_idx];
                # use processing pool
                pool = multiprocessing.Pool(processes=4);
                for QP_idx in range(4):
                    QP = QP_list[QP_idx];
                    if map_file_name[-len(map_file_ext):] == map_file_ext:
                        print(r"process _map.yuv file:"+map_file_name);
                        output_name = map_file_name[:-len(map_file_ext)]+str(QP)+avs_file_ext;
                        ReconFile_name = file_name_all[file_name_idx_rec][map_file_name_idx][QP_idx];
                        if os.path.isfile(os.path.join(AVS_path, output_name)):
                            print(os.path.join(AVS_path, output_name) + r" exist, delete it");
                            os.remove(os.path.join(MAP_path, output_name));

                        command = os.path.join(RD_bin, ENCODER) + r" " \
                                  + encoder_para \
                                  + r" -p InputFile="+map_file_name \
                                  + r" -p SourceWidth="+map_yuv_width \
                                  + r" -p SourceHeight="+map_yuv_height \
                                  + r" -p OutputFile="+os.path.join(root_path, os.path.join(AVS_path, output_name));
                        os.system(r"echo RD start...");

                        os.system(r"echo QP="+str(QP));
                        QP_I = QP;
                        QP_P = QP+1;
                        QP_B = QP+4;
                        RD_log = r"RD_"+str(QP)+r'.log';
                        command += r" -p QPIFrame="+str(QP_I) \
                                   + r" -p QPPFrame="+str(QP_P) \
                                   + r" -p QPBFrame="+str(QP_B) \
                                   + r" -p ReconFile="+ReconFile_name \
                                   + '>>'+RD_log;
                        pool.apply_async(os.system, (command,));
                        #os.system(command);
                pool.close();
                pool.join();

    if SPSNR_flag:
        #read the yaw roll pitch and write to cfg

        #decoder
        print(r"RD ldocod: decode the seq...")
        os.chdir(os.path.join(root_path, VRM_path));
        VRM_path_all = os.path.join(root_path, os.path.join(VRM_path, VRM));
        command = r" -w 3840 -h 1920 -x 3840 -y 1920 -N 30 -M psnr -s sphere_655362.txt -F eqar -f eqar -c encoder_ra.cfg";

        src_file_name_idx = file_name_idx[r"src"];
        map_file_name_idx = file_name_idx[r"map"];
        rec_file_name_idx = file_name_idx[r"rec"];
        for file_idx in range(len(file_name_all[map_file_name_idx])):
            file_name = file_name_all[src_file_name_idx][file_idx];
            pool = multiprocessing.Pool(processes=4);
            for rec_file_name_idx_1 in range(4):
                rec_file_name = file_name_all[rec_file_name_idx][file_idx][rec_file_name_idx_1];
                QP = QP_list[rec_file_name_idx_1];
                VRM_log = r"VRM_PSNR_" + str(QP) + r".log";
                input_file_name = os.path.join(root_path, os.path.join(SEQ_path, file_name));
                output_file_name = os.path.join(root_path, os.path.join(AVS_path, rec_file_name));

                SPSNR_command = VRM_path_all \
                                + r" -i " + input_file_name \
                                + r" -o " + output_file_name \
                                + command + r">>" + VRM_log;

                pool.apply_async(os.system, (SPSNR_command,));

            pool.close();
            pool.join();

        print(r"RD end");

if __name__ == "__main__":
    root_path = '';
    flag = 0;
    options, remainder = getopt.getopt(sys.argv[1:], 'r', ['root']);
    for opt, arg in options:
        if opt in ('-r', '--root'):
            root_path = arg;
        elif opt in (r'-f', r'-flag'):
            flag = int(arg);
        else:
            help_info();

    flag = 0b111;
    process_all(root_path, flag);
